 # This code generate scores per breast level
# for each breast combine MLO and CC views
import sys
# here is my py-faster-rcnn installation
sys.path.append('py-faster-rcnn/tools/')
sys.path.append('py-faster-rcnn/caffe-fast-rcnn/python/')
import _init_paths
from fast_rcnn.config import cfg
from fast_rcnn.test import im_detect
from fast_rcnn.nms_wrapper import nms
import numpy as np

import caffe, cv2
import os

import csv

CLASSES = ('__background__','benign','malignant')

def load_net(ptxt, w, device=0):
    """Load model."""
    caffe.set_mode_gpu()
    caffe.set_device(device)
    cfg.GPU_ID = device
    net = caffe.Net(ptxt, w, caffe.TEST)
    return net

def detect(net, im,  NMS_THRESH = 0.1, CLASSES = ('__background__','benign','malignant')):
    """Detect cancer."""
    cfg.TEST.HAS_RPN = True
    scores, boxes = im_detect(net, im)
    dets_all_cls = []
    for cls_ind, cls in enumerate(CLASSES[1:]):
        cls_ind += 1 # because we skipped background
        cls_boxes = boxes[:, 4*cls_ind:4*(cls_ind + 1)]
        cls_scores = scores[:, cls_ind]
        dets = np.hstack((cls_boxes,
                          cls_scores[:, np.newaxis])).astype(np.float32)
        keep = nms(dets, NMS_THRESH)
        dets = dets[keep, :]
        dets_all_cls.append(dets)
    return dets_all_cls


def vis_detections(im, dets_all_cls, out_filename, CONF_THRESH = [0.5,0.5],
                   CLASSES = ('__background__','benign','malignant'),vis_gt = False):
    import matplotlib.pyplot as plt
    plt.interactive(False)
    """Draw boxes around detected cancer."""
    fig,ax=plt.subplots(figsize=(8,10))
    ax.imshow(im,cmap='gray_r')
    # if vis_gt == True:
    #     birads = out_filename.split('_')[-2]
    #     pathology = out_filename.split('_')[-1][0]
    #     ax.text(20, 50, '{:s} {:s}'.format('Birads : ', birads),
    #                 fontsize=14, color='yellow')
    #     ax.text(120, 100, '{:s} {:s}'.format('Pathology : ', pathology),
    #                 fontsize=14, color='yellow')

    for ii, cls in enumerate(CLASSES[1:]):

        dets = dets_all_cls[ii]
        inds = np.where(dets[:, -1] >= CONF_THRESH[ii])[0]
        for i in inds:
            bbox = dets[i, :4]
            score = dets[i, -1]
            box_color = (0.9,0.1,0.1) #
            ax.add_patch(
                plt.Rectangle((bbox[0], bbox[1]),
                              bbox[2] - bbox[0],
                              bbox[3] - bbox[1], fill=False, linestyle ='dashed',
                              edgecolor=box_color, linewidth=3))
            ax.text(bbox[0], bbox[1] - 2,
                        '{:s} {:.2f}'.format(CLASSES[ii+1], score),
                        bbox=dict(facecolor='blue', alpha=0.5),
                        fontsize=14, color='white')

    plt.axis('off')
    plt.tight_layout()
    # plt.show()
    plt.savefig(out_filename)
    plt.close()

def plot_roc(y_pred, y_gt,out_filename = 'roc.png'):
    import matplotlib.pyplot as plt
    plt.interactive(False)
    from sklearn.metrics import roc_curve, auc
    fpr, tpr, thresholds = roc_curve(y_gt, y_pred)
    roc_auc = auc(fpr, tpr)
    plt.plot(fpr, tpr, lw=3, label='ROC (area = %0.2f)' % roc_auc)
    plt.xlim([0, 1])
    plt.ylim([0, 1])
    plt.grid(True,linestyle = ":")
    plt.xlabel('False Positive Rate (1 - Specificity)',fontsize=15)
    plt.ylabel('True Positive Rate (Sensitivity)',fontsize=15)
    plt.title('Receiver operating characteristic example',fontsize=15)
    plt.legend(loc="lower right")
    plt.show()
    plt.tight_layout()
    plt.savefig(out_filename)
    plt.close()


if __name__ == '__main__':
    import pandas as pd

    # faster rcnn configs we need to change
    cfg.TEST.HAS_RPN = True  # Use RPN for proposals
    cfg.TEST.SCALES=(1000,)  # change scales
    cfg.TEST.MAX_SIZE= 2000 # change scales
    # define root path
    root_path = '/home/he/py-faster-rcnn1/py-faster-rcnn/'



    # define other pathes
    im_dir = os.path.join(root_path, 'data', 'demo_ddsm_longside=2000')
    out_dir = os.path.join(root_path,'results','detections6')
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    
    # get all images in im_dir
    im_names = os.listdir(im_dir)

    # loop evaluate per case, generate a score
    net = load_net(ptxt=root_path + 'models/pascal_voc/VGG16/faster_rcnn_end2end/test.prototxt',
                       w=root_path + '/data/faster_rcnn_models/vgg16_faster_rcnn_iter_40000_longside=2000.caffemodel')

    scores = np.zeros(len(im_names)) # malignant scores for each image
    gt = np.zeros(len(im_names))
    #fcsv = open('/home/he/py-faster-rcnn1/py-faster-rcnn/scores.csv','w')
    #writer = csv.writer(fcsv)
    for index, im_name in enumerate(im_names):
        print('Detecting image {}'.format(im_name))
        filename = os.path.join(im_dir,im_name)
        im = cv2.imread(filename)
        detections= detect(net, im)

        vis_detections(im, detections, out_dir+'/'+
                       im_name, CONF_THRESH = [0.5,0.5]) # filter out all detections with confidence lower than 0.4

        # here detections is a 4xndx5 array. The first dimension is the class. 0 means benign, 1 means malignant
        # The second dimension is the nd detections of the class sorting by descending order according to the confidence,
        # nd is the number of detections. 0 means the first element, which means the highest confidence detection.
        # The third dimension has 5 elements, showing the position of the bounding box (4 elements), and the last elements
        # as the confidence score.

        # gt_tmp = int(filename[-5])

	if filename[-5] == 'T':
	    gt_tmp = 1
	else:
	    gt_tmp = 0
	
        if detections[1].size != 0:
            scores[index] = detections[1][0][-1]
            gt[index] = gt_tmp
	#writer.writerow([filename,detections[1][0][-1]])
    #fcsv.close()


    # save predictions as a csv file
    np.savetxt(out_dir+'/scores.csv', scores, delimiter=",")
    np.savetxt(out_dir+'/gt.csv', gt, delimiter=",")
    # Plot ROC curve, calculate AUC score
    
    #scores = np.genfromtxt('/home/he/py-faster-rcnn1/py-faster-rcnn/results/detections5/scores.csv', delimiter=',')

    #gt = np.genfromtxt('/home/he/py-faster-rcnn1/py-faster-rcnn/results/detections5/gt.csv', delimiter=',')
  
    #print(scores)
    #print(gt)

    plot_roc(scores,gt,out_filename =out_dir+'/roc.png')

